from airflow.models import DagRun, TaskInstance
from airflow.models.idata_task import OperatorType, Task
from airflow.models.idata_workflow_task import WorkFlowTask


class DagUtil(object):

    @staticmethod
    def recursion_clear_history(dag_id, execution_date, session=None):
        wid = dag_id.split("-")[-1]
        session.query(DagRun).filter(DagRun.execution_date == execution_date,
                                     DagRun.dag_id == dag_id).delete()
        session.query(TaskInstance).filter(TaskInstance.dag_id == dag_id,
                                           TaskInstance.execution_date == execution_date).delete()
        wfts = WorkFlowTask.query_workflow_task(wid, session)
        for wft in wfts:
            task = Task.get_task(wft.task_id, session)
            if task.operator == OperatorType.DAGRUN:
                task_param = task.param
                child_wid = task_param.get("workflow_id", None)
                if child_wid:
                    child_dag_id = 'dag-{}'.format(child_wid)
                    DagUtil.recursion_clear_history(child_dag_id, execution_date, session)





